use bytes::{Bytes, BytesMut, Buf, BufMut};
use std::io::{self, Read, Write};
use std::net::TcpStream;

use native_tls::{TlsConnector, TlsStream};

use crate::error::{Result, VerticaError};

/// A buffered reader for network communication
pub struct BufferedReader<R> {
    reader: R,
    buffer: BytesMut,
    capacity: usize,
}
impl<R: Read> BufferedReader<R> {
    /// 创建一个新的 `BufferedReader` 实例
    ///
    /// # 参数
    /// - `reader`: 实现了 `Read`  trait 的读取器
    /// - `capacity`: 缓冲区的初始容量
    pub fn new(reader: R, capacity: usize) -> Self {
        Self {
            reader,
            buffer: BytesMut::with_capacity(capacity),
            capacity,
        }
    }

    /// 从缓冲区或底层读取器中读取指定长度的字节
    ///
    /// 如果缓冲区中的字节数不足，则从底层读取器中读取剩余所需的字节
    ///
    /// # 参数
    /// - `len`: 需要读取的字节数
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时包含读取到的 `Bytes`，失败时包含错误信息
    pub fn read_exact(&mut self, len: usize) -> Result<Bytes> {
        if self.buffer.len() < len {
            // 计算还需要读取的字节数
            let needed = len - self.buffer.len();
            let mut temp = vec![0u8; needed];
            // 从底层读取器中精确读取所需字节数
            self.reader.read_exact(&mut temp)?;
            // 将读取到的字节添加到缓冲区
            self.buffer.extend_from_slice(&temp);
        }

        // 从缓冲区中分割出指定长度的字节并冻结
        Ok(self.buffer.split_to(len).freeze())
    }

    /// 从缓冲区或底层读取器中读取字节，直到遇到指定的分隔符
    ///
    /// # 参数
    /// - `delimiter`: 分隔符字节
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时包含从开始到分隔符的 `Bytes`，失败时包含错误信息
    pub fn read_until(&mut self, delimiter: u8) -> Result<Bytes> {
        // 检查缓冲区中是否已存在分隔符
        if let Some(pos) = self.buffer.iter().position(|&b| b == delimiter) {
            return Ok(self.buffer.split_to(pos + 1).freeze());
        }

        let mut temp = vec![0u8; 1024];
        loop {
            // 从底层读取器中读取数据
            let bytes_read = self.reader.read(&mut temp)?;
            if bytes_read == 0 {
                // 读取到文件末尾，连接已关闭
                return Err(VerticaError::Io(std::io::Error::new(io::ErrorKind::Other, "Connection closed")));
            }
            // 将读取到的字节添加到缓冲区
            self.buffer.extend_from_slice(&temp[..bytes_read]);

            // 再次检查缓冲区中是否已存在分隔符
            if let Some(pos) = self.buffer.iter().position(|&b| b == delimiter) {
                return Ok(self.buffer.split_to(pos + 1).freeze());
            }
        }
    }

    /// 从缓冲区或底层读取器中读取一个完整的消息
    ///
    /// 消息格式为：1 字节消息类型 + 4 字节消息长度 + 消息内容
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时包含消息内容的 `Bytes`，失败时包含错误信息
    pub fn read_message(&mut self) -> Result<Bytes> {
        // 确保缓冲区中至少有 5 字节（1 字节消息类型 + 4 字节消息长度）
        if self.buffer.len() < 5 {
            let mut temp = vec![0u8; 1024];
            let bytes_read = self.reader.read(&mut temp)?;
            if bytes_read == 0 {
                // 读取到文件末尾，连接已关闭
                return Err(VerticaError::Io(std::io::Error::new(io::ErrorKind::UnexpectedEof, "Connection closed")));
            }
            self.buffer.extend_from_slice(&temp[..bytes_read]);
        }

        // 再次检查缓冲区中是否有足够的字节
        if self.buffer.len() < 5 {
            return Err(VerticaError::Io(io::Error::new(io::ErrorKind::UnexpectedEof, "Connection closed")));
        }

        // 获取消息类型（未使用）
        let _message_type = self.buffer[0];
        // 从缓冲区中解析出消息长度
        let message_length = i32::from_be_bytes([self.buffer[1], self.buffer[2], self.buffer[3], self.buffer[4]]);
        
        if message_length < 4 {
            // 消息长度无效
            return Err(VerticaError::Io(io::Error::new(io::ErrorKind::InvalidData, "Invalid message length")));
        }

        // 计算消息总长度（消息类型 + 消息长度 + 消息内容）
        let total_length = 5 + message_length as usize - 4;
        
        // 循环读取数据，直到缓冲区中有足够的字节
        while self.buffer.len() < total_length {
            let mut temp = vec![0u8; 1024];
            let bytes_read = self.reader.read(&mut temp)?;
            if bytes_read == 0 {
                // 读取到文件末尾，连接已关闭
                return Err(VerticaError::Io(std::io::Error::new(io::ErrorKind::UnexpectedEof, "Connection closed")));
            }
            self.buffer.extend_from_slice(&temp[..bytes_read]);
        }

        // 从缓冲区中分割出完整的消息
        let mut data = self.buffer.split_to(total_length);
        // 跳过消息类型字节
        let _message_type = data.get_u8(); 
        // 跳过消息长度字节
        let _length = data.get_i32(); 
        // 将剩余的消息内容冻结并返回
        Ok(data.freeze())
    }

    /// 获取当前缓冲区的引用
    ///
    /// # 返回
    /// 返回当前缓冲区的引用
    pub fn buffer(&self) -> &BytesMut {
        &self.buffer
    }

    /// 清空当前缓冲区
    pub fn clear_buffer(&mut self) {
        self.buffer.clear();
    }
}

/// 用于网络通信的缓冲写入器
pub struct BufferedWriter<W> {
    writer: W,
    buffer: BytesMut,
    capacity: usize,
}

impl<W: Write> BufferedWriter<W> {
    /// 创建一个新的 `BufferedWriter` 实例
    ///
    /// # 参数
    /// - `writer`: 实现了 `Write`  trait 的写入器
    /// - `capacity`: 缓冲区的初始容量
    pub fn new(writer: W, capacity: usize) -> Self {
        Self {
            writer,
            buffer: BytesMut::with_capacity(capacity),
            capacity,
        }
    }

    /// 将数据写入缓冲区，如果缓冲区已满则刷新缓冲区
    ///
    /// # 参数
    /// - `data`: 需要写入的字节切片
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn write(&mut self, data: &[u8]) -> Result<()> {
        // 将数据添加到缓冲区
        self.buffer.extend_from_slice(data);
        
        if self.buffer.len() >= self.capacity {
            // 缓冲区已满，刷新缓冲区
            self.flush()?;
        }

        Ok(())
    }

    /// 写入一个完整的消息到缓冲区
    ///
    /// 消息格式为：1 字节消息类型 + 4 字节消息长度 + 消息内容
    ///
    /// # 参数
    /// - `message_type`: 消息类型字节
    /// - `payload`: 消息内容
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn write_message(&mut self, message_type: u8, payload: &[u8]) -> Result<()> {
        // 计算消息总长度（消息类型 + 消息长度 + 消息内容）
        let total_length = 1 + 4 + payload.len();
        
        if total_length > self.capacity {
            // 消息过长，无法写入缓冲区
            return Err(VerticaError::Io(io::Error::new(io::ErrorKind::Other, "Message too large")));
        }

        if self.buffer.len() + total_length > self.capacity {
            // 缓冲区空间不足，刷新缓冲区
            self.flush()?;
        }

        // 写入消息类型
        self.buffer.put_u8(message_type);
        // 写入消息长度
        self.buffer.put_i32((payload.len() + 4) as i32);
        // 写入消息内容
        self.buffer.extend_from_slice(payload);

        Ok(())
    }

    /// 刷新缓冲区，将缓冲区中的数据写入底层写入器
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn flush(&mut self) -> Result<()> {
        if !self.buffer.is_empty() {
            // 将缓冲区中的数据写入底层写入器
            self.writer.write_all(&self.buffer)?;
            // 清空缓冲区
            self.buffer.clear();
        }
        Ok(())
    }

    /// 获取当前缓冲区的引用
    ///
    /// # 返回
    /// 返回当前缓冲区的引用
    pub fn buffer(&self) -> &BytesMut {
        &self.buffer
    }
}
/// Network connection wrapper
pub struct NetworkStream {
    stream: TcpStream,
    reader: BufferedReader<TcpStream>,
    writer: BufferedWriter<TcpStream>,
}
impl NetworkStream {
    /// 创建一个新的 `NetworkStream` 实例
    ///
    /// 该方法会克隆传入的 `TcpStream` 以分别用于读取和写入操作，
    /// 并初始化一个容量为 8192 的 `BufferedReader` 和 `BufferedWriter`。
    ///
    /// # 参数
    /// - `stream`: 基础的 `TcpStream` 实例
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时包含新创建的 `NetworkStream` 实例，失败时包含错误信息
    pub fn new(stream: TcpStream) -> Result<Self> {
        // 克隆流以用于读取操作
        let reader_stream = stream.try_clone()?;
        // 克隆流以用于写入操作
        let writer_stream = stream.try_clone()?;
        
        Ok(Self {
            stream,
            reader: BufferedReader::new(reader_stream, 8192),
            writer: BufferedWriter::new(writer_stream, 8192),
        })
    }

    /// 从网络流中读取一个完整的消息
    ///
    /// 消息格式为：1 字节消息类型 + 4 字节消息长度 + 消息内容
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时包含消息内容的 `Bytes`，失败时包含错误信息
    pub fn read_message(&mut self) -> Result<Bytes> {
        self.reader.read_message()
    }

    /// 写入一个完整的消息到网络流
    ///
    /// 消息格式为：1 字节消息类型 + 4 字节消息长度 + 消息内容
    ///
    /// # 参数
    /// - `message_type`: 消息类型字节
    /// - `payload`: 消息内容
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn write_message(&mut self, message_type: u8, payload: &[u8]) -> Result<()> {
        self.writer.write_message(message_type, payload)
    }

    /// 刷新写入缓冲区，将缓冲区中的数据写入底层流
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn flush(&mut self) -> Result<()> {
        self.writer.flush()
    }

    /// 关闭网络连接
    ///
    /// 该方法会先刷新写入缓冲区，确保所有数据都已写入底层流，
    /// 然后关闭连接的读写通道。
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn shutdown(&mut self) -> Result<()> {
        // 确保所有缓冲数据都已写入底层流
        self.flush()?;
        // 关闭连接的读写通道
        self.stream.shutdown(std::net::Shutdown::Both)?;
        Ok(())
    }
}

/// Network connection wrapper
pub enum ConnectionStream {
    /// Plain TCP connection
    Plain(TcpStream),
    /// TLS encrypted connection
    Tls(TlsStream<TcpStream>),
}

impl ConnectionStream {
    /// 创建一个新的普通 TCP 连接
    ///
    /// # 参数
    /// - `stream`: 基础的 `TcpStream` 实例
    ///
    /// # 返回
    /// 返回一个新的 `ConnectionStream::Plain` 实例
    pub fn plain(stream: TcpStream) -> Self {
        ConnectionStream::Plain(stream)
    }

    /// 创建一个新的 TLS 加密连接
    ///
    /// # 参数
    /// - `stream`: 基础的 `TlsStream<TcpStream>` 实例
    ///
    /// # 返回
    /// 返回一个新的 `ConnectionStream::Tls` 实例
    pub fn tls(stream: TlsStream<TcpStream>) -> Self {
        ConnectionStream::Tls(stream)
    }

    /// 将普通 TCP 连接升级为 TLS 加密连接
    ///
    /// 如果当前连接已经是 TLS 加密连接，则直接返回 `Ok(())`。
    ///
    /// # 参数
    /// - `connector`: `TlsConnector` 实例，用于建立 TLS 连接
    /// - `domain`: 要连接的域名
    ///
    /// # 返回
    /// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
    pub fn upgrade_tls(&mut self, connector: &TlsConnector, domain: &str) -> Result<()> {
        match self {
            ConnectionStream::Plain(stream) => {
                // 使用 TlsConnector 连接到指定域名
                let tls_stream = connector.connect(domain, stream.try_clone()?)
                    .map_err(|e| match e {
                        native_tls::HandshakeError::Failure(e) => VerticaError::Connection(e.to_string()),
                        native_tls::HandshakeError::WouldBlock(_) => VerticaError::Connection("TLS handshake would block".to_string()),
                    })?;
                // 将当前连接替换为 TLS 连接
                *self = ConnectionStream::Tls(tls_stream);
                Ok(())
            }
            ConnectionStream::Tls(_) => Ok(()), // Already TLS
        }
    }
}

/// 为 `ConnectionStream` 实现 `Read` trait，使其支持从流中读取数据
impl Read for ConnectionStream {
    /// 从 `ConnectionStream` 中读取数据到指定缓冲区
    ///
    /// # 参数
    /// - `buf`: 用于存储读取数据的缓冲区
    ///
    /// # 返回
    /// 返回一个 `io::Result`，成功时包含实际读取的字节数，失败时包含错误信息
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        match self {
            ConnectionStream::Plain(stream) => stream.read(buf),
            ConnectionStream::Tls(stream) => stream.read(buf),
        }
    }
}

/// 为 `ConnectionStream` 实现 `Write` trait，使其支持向流中写入数据
impl Write for ConnectionStream {
    /// 将指定缓冲区中的数据写入 `ConnectionStream`
    ///
    /// # 参数
    /// - `buf`: 包含要写入数据的缓冲区
    ///
    /// # 返回
    /// 返回一个 `io::Result`，成功时包含实际写入的字节数，失败时包含错误信息
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            ConnectionStream::Plain(stream) => stream.write(buf),
            ConnectionStream::Tls(stream) => stream.write(buf),
        }
    }

    /// 刷新 `ConnectionStream`，确保所有缓冲数据都已写入底层流
    ///
    /// # 返回
    /// 返回一个 `io::Result`，成功时为 `Ok(())`，失败时包含错误信息
    fn flush(&mut self) -> io::Result<()> {
        match self {
            ConnectionStream::Plain(stream) => stream.flush(),
            ConnectionStream::Tls(stream) => stream.flush(),
        }
    }
}

/// 从流中精确读取指定长度的字节
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
/// - `buf`: 用于存储读取数据的缓冲区
///
/// # 返回
/// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
pub fn read_exact(stream: &mut ConnectionStream, buf: &mut [u8]) -> Result<()> {
    stream.read_exact(buf).map_err(|e| VerticaError::Io(e))
}

/// 将字节数据写入流
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
/// - `buf`: 包含要写入数据的缓冲区
///
/// # 返回
/// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
pub fn write_all(stream: &mut ConnectionStream, buf: &[u8]) -> Result<()> {
    stream.write_all(buf).map_err(|e| VerticaError::Io(e))
}

/// 从流中读取一个字节
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
///
/// # 返回
/// 返回一个 `Result`，成功时包含读取到的字节，失败时包含错误信息
pub fn read_u8(stream: &mut ConnectionStream) -> Result<u8> {
    let mut buf = [0u8; 1];
    // 精确读取一个字节
    read_exact(stream, &mut buf)?;
    Ok(buf[0])
}

/// 从流中读取一个 32 位大端序无符号整数
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
///
/// # 返回
/// 返回一个 `Result`，成功时包含读取到的 32 位无符号整数，失败时包含错误信息
pub fn read_u32(stream: &mut ConnectionStream) -> Result<u32> {
    let mut buf = [0u8; 4];
    // 精确读取 4 个字节
    read_exact(stream, &mut buf)?;
    Ok(u32::from_be_bytes(buf))
}

/// 从流中读取一个 32 位大端序有符号整数
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
///
/// # 返回
/// 返回一个 `Result`，成功时包含读取到的 32 位有符号整数，失败时包含错误信息
pub fn read_i32(stream: &mut ConnectionStream) -> Result<i32> {
    let mut buf = [0u8; 4];
    // 精确读取 4 个字节
    read_exact(stream, &mut buf)?;
    Ok(i32::from_be_bytes(buf))
}

/// 将一个 32 位大端序无符号整数写入流
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
/// - `value`: 要写入的 32 位无符号整数
///
/// # 返回
/// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
pub fn write_u32(stream: &mut ConnectionStream, value: u32) -> Result<()> {
    let buf = value.to_be_bytes();
    write_all(stream, &buf)
}

/// 将一个 32 位大端序有符号整数写入流
///
/// # 参数
/// - `stream`: 可变引用的 `ConnectionStream` 实例
/// - `value`: 要写入的 32 位有符号整数
///
/// # 返回
/// 返回一个 `Result`，成功时为 `Ok(())`，失败时包含错误信息
pub fn write_i32(stream: &mut ConnectionStream, value: i32) -> Result<()> {
    let buf = value.to_be_bytes();
    write_all(stream, &buf)
}

#[cfg(test)]
mod tests {
    use super::*;

    /// 测试 `ConnectionStream` 的创建功能
    #[test]
    fn test_connection_stream_creation() {
        // Test creation functions
        // 尝试连接到本地端口 0，如果失败则尝试端口 1 和 2 作为测试备用方案
        let tcp_stream = TcpStream::connect("127.0.0.1:0").unwrap_or_else(|_| {
            // Fallback for testing
            TcpStream::connect_timeout(&"127.0.0.1:1".parse().unwrap(), std::time::Duration::from_millis(1)).unwrap_or_else(|_| {
                TcpStream::connect_timeout(&"127.0.0.1:2".parse().unwrap(), std::time::Duration::from_millis(1)).unwrap()
            })
        });
        
        // 创建一个普通 TCP 连接
        let plain = ConnectionStream::plain(tcp_stream.try_clone().unwrap());
        // 验证创建的连接类型是否为普通 TCP 连接
        assert!(matches!(plain, ConnectionStream::Plain(_)));
    }
}